package main

import (
	"fmt"
	"math"
	"math/rand/v2"
	"sort"
	"time"
)

// LoadDataPoint 表示一个时间点的负荷数据
type LoadDataPoint struct {
	Timestamp   time.Time
	Value       float64 // 负荷值 (kW)
	Temperature float64 // 温度影响
}

// LoadPredictor 负荷预测器结构体
type LoadPredictor struct {
	Data         []LoadDataPoint // 历史负荷数据
	SeasonalCoef map[int]float64 // 季节性系数 (小时 -> 系数)
}

// NewLoadPredictor 创建新的负荷预测器
func NewLoadPredictor() *LoadPredictor {
	return &LoadPredictor{
		SeasonalCoef: make(map[int]float64),
	}
}

// AddData 添加历史数据点
func (lp *LoadPredictor) AddData(timestamp time.Time, value float64) {
	lp.Data = append(lp.Data, LoadDataPoint{Timestamp: timestamp, Value: value})
}

// Train 训练模型，计算季节性系数
func (lp *LoadPredictor) Train() {
	// 按小时分组数据
	hourlyData := make(map[int][]float64)
	for _, dp := range lp.Data {
		hour := dp.Timestamp.Hour()
		hourlyData[hour] = append(hourlyData[hour], dp.Value)
	}

	// 计算每小时的平均负荷
	hourlyAvg := make(map[int]float64)
	for hour, values := range hourlyData {
		sum := 0.0
		for _, v := range values {
			sum += v
		}
		hourlyAvg[hour] = sum / float64(len(values))
	}

	// 计算全局平均负荷
	totalAvg := 0.0
	for _, avg := range hourlyAvg {
		totalAvg += avg
	}
	totalAvg /= float64(len(hourlyAvg))

	// 计算季节性系数 (每小时负荷 / 全局平均负荷)
	for hour, avg := range hourlyAvg {
		lp.SeasonalCoef[hour] = avg / totalAvg
	}
}

// Predict 预测未来指定时间的负荷
func (lp *LoadPredictor) Predict(timestamp time.Time, windowSize int) float64 {
	if len(lp.Data) == 0 {
		return 0.0
	}

	// 获取季节性系数
	hour := timestamp.Hour()
	seasonalFactor := lp.SeasonalCoef[hour]
	if seasonalFactor == 0 {
		// 如果没有该小时的数据，使用最近的小时系数
		seasonalFactor = lp.getClosestSeasonalFactor(hour)
	}

	// 计算加权移动平均
	baseLoad := lp.weightedMovingAverage(windowSize)
	// 应用季节性调整
	return baseLoad * seasonalFactor
}

// weightedMovingAverage 计算加权移动平均
func (lp *LoadPredictor) weightedMovingAverage(windowSize int) float64 {
	n := len(lp.Data)
	if n < windowSize {
		windowSize = n
	}

	// 获取最近的 windowSize 个数据点
	recentData := lp.Data[n-windowSize:]

	// 计算加权平均（最近的权重更高）
	totalWeight := 0.0
	weightedSum := 0.0
	for i, dp := range recentData {
		weight := float64(i + 1) // 线性权重
		totalWeight += weight
		weightedSum += dp.Value * weight
	}

	return weightedSum / totalWeight
}

// getClosestSeasonalFactor 获取最近的小时系数
func (lp *LoadPredictor) getClosestSeasonalFactor(hour int) float64 {
	// 获取所有有系数的小时
	hours := make([]int, 0, len(lp.SeasonalCoef))
	for h := range lp.SeasonalCoef {
		hours = append(hours, h)
	}
	sort.Ints(hours)

	// 找到最接近的小时
	minDiff := 24
	closestHour := hour
	for _, h := range hours {
		diff := int(math.Abs(float64(h - hour)))
		if diff < minDiff {
			minDiff = diff
			closestHour = h
		}
	}

	return lp.SeasonalCoef[closestHour]
}

// 示例使用
func main() {
	// 创建预测器
	predictor := NewLoadPredictor()

	// 生成模拟历史数据 (每小时一个点)
	now := time.Now().Add(-time.Hour * 24 * 7) // 一周前
	for i := 0; i < 168; i++ {                 // 7天 * 24小时
		// 模拟负荷：基值 + 日周期 + 随机波动
		hour := now.Hour()
		baseLoad := 500.0
		hourlyVariation := math.Sin(float64(hour)/24.0*2*math.Pi) * 200.0
		randomVariation := rand.Float64() * 50.0

		value := baseLoad + hourlyVariation + randomVariation
		predictor.AddData(now, value)
		now = now.Add(time.Hour)
	}

	// 训练模型
	predictor.Train()

	// 预测未来24小时的负荷
	fmt.Println("未来24小时负荷预测:")
	startTime := time.Now().Add(time.Hour)
	for i := 0; i < 24; i++ {
		predicted := predictor.Predict(startTime, 24) // 使用最近24个点计算移动平均
		fmt.Printf("%s: %.2f kW\n", startTime.Format("15:04"), predicted)
		startTime = startTime.Add(time.Hour)
	}
}
